from litellm import Router
import os
import asyncio
import litellm
from custom_router import CustomRoutingStrategy
import random
import json
from typing import Dict, List, Optional, Union
from litellm.caching import Cache

# 设置API密钥和基础URL  
os.environ["OPENAI_API_KEY"] = "sk-hqNXyGi0PFoZOI4mEc8a6696A04148Ac921a97C866BbB81d"
os.environ["OPENAI_API_BASE"] = "https://api.zyfan.work/v1"
os.environ["HUAWEI_API_KEY"] = "rxCsUK0v19j0yZoaduUDvliyDwcouqpu_fovkSwyI9GGK5qaL4z2SITnvAGx5bj5PsuT3lHCto5DB_NtR5xoQyA"

# 设置Redis环境变量
os.environ["REDIS_HOST"] = "redis-14342.c52.us-east-1-4.ec2.redns.redis-cloud.com"
os.environ["REDIS_PORT"] = "14342"
os.environ["REDIS_PASSWORD"] = "UGwqvLGG1aoCSjXnXFDt26cC8TryQ3ER"

# LangSmith配置
os.environ["LANGSMITH_API_KEY"] = "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"
os.environ["LANGSMITH_PROJECT"] = "litellm"
os.environ["LANGSMITH_DEFAULT_RUN_NAME"] = "LLMRun"

litellm.success_callback = ["langsmith"]

model_list = [
    # {
    #     "model_name": "deepseek-big-R1",
    #     "litellm_params": {
    #         "model": "deepseek/DeepSeek-R1",
    #         "api_key": os.getenv("HUAWEI_API_KEY"),
    #         "api_base": os.getenv("https://infer-modelarts-cn-southwest-2.modelarts-infer.com/v1/infers/952e4f88-ef93-4398-ae8d-af37f63f0d8e/v1"),
    #     },
    #     "model_info": {
    #         "model_from": "huawei_DeepSeek-R1-671B-32K",
    #         "first_token_time": "2",
    #         "token_per_second": "56",
    #         "cost_per_token": "0.12",
    #     }
    # },
    # {
    #     "model_name": "deepseek-big-R1",
    #     "litellm_params": {
    #         "model": "openai/DeepSeek-R1",
    #         "api_key": os.getenv("HUAWEI_API_KEY"),
    #         "api_base": os.getenv("https://infer-modelarts-cn-southwest-2.modelarts-infer.com/v1/infers/861b6827-e5ef-4fa6-90d2-5fd1b2975882/v1"),
    #     },
    #     "model_info": {
    #         "model_from": "huawei_DeepSeek-R1-671B-8K",
    #         "first_token_time": "2",
    #         "token_per_second": "56",
    #         "cost_per_token": "0.12",
    #     }
    # },
    # {
    #     "model_name": "deepseek-big-R1",
    #     "litellm_params": {
    #         "model": "openai/DeepSeek-R1",
    #         "api_key": os.getenv("HUAWEI_API_KEY"),
    #         "api_base": os.getenv("https://infer-modelarts-cn-southwest-2.modelarts-infer.com/v1/infers/c3cfa9e2-40c9-485f-a747-caae405296ef/v1"),
    #     },
    #     "model_info": {
    #         "model_from": "huawei_DeepSeek-R1-671B-4K",
    #         "first_token_time": "2",
    #         "token_per_second": "56",
    #         "cost_per_token": "0.12",
    #     }
    # },
    # {
    #     "model_name": "deepseek-big-V3",
    #     "litellm_params": {
    #         "model": "openai/DeepSeek-V3",
    #         "api_key": os.getenv("HUAWEI_API_KEY"),
    #         "api_base": os.getenv("https://infer-modelarts-cn-southwest-2.modelarts-infer.com/v1/infers/fd53915b-8935-48fe-be70-449d76c0fc87/v1"),
    #     },
    #     "model_info": {
    #         "model_from": "huawei_DeepSeek-V3-671B-32K",
    #         "first_token_time": "2",
    #         "token_per_second": "56",
    #         "cost_per_token": "0.12",
    #     }
    # },
    # {
    #     "model_name": "deepseek-big-V3",
    #     "litellm_params": {
    #         "model": "openai/DeepSeek-V3",
    #         "api_key": os.getenv("HUAWEI_API_KEY"),
    #         "api_base": os.getenv("https://infer-modelarts-cn-southwest-2.modelarts-infer.com/v1/infers/707c01c8-517c-46ca-827a-d0b21c71b074/v1"),
    #     },
    #     "model_info": {
    #         "model_from": "huawei_DeepSeek-V3-671B-8K",
    #         "first_token_time": "2",
    #         "token_per_second": "56",
    #         "cost_per_token": "0.12",
    #     }
    # },
    # {
    #     "model_name": "deepseek-big-V3",
    #     "litellm_params": {
    #         "model": "openai/DeepSeek-V3",
    #         "api_key": os.getenv("HUAWEI_API_KEY"),
    #         "api_base": os.getenv("https://infer-modelarts-cn-southwest-2.modelarts-infer.com/v1/infers/f354eacc-a2c5-43b4-a785-e5aadca988b3/v1"),
    #     },
    #     "model_info": {
    #         "model_from": "huawei_DeepSeek-V3-671B-4K",
    #         "first_token_time": "2",
    #         "token_per_second": "56",
    #         "cost_per_token": "0.12",
    #     }
    # }  
    {
        "model_name": "deepseek-big",
        "litellm_params": {
            "model": "openai/gpt-4o",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "2",
            "token_per_second": "56",
            "cost_per_token": "0.12",
        }
    },
    {
        "model_name": "deepseek-big",
        "litellm_params": {
            "model": "openai/gpt-4o-2024-08-06",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "3",
            "token_per_second": "34",
            "cost_per_token": "0.11",
        }
    },
    {
        "model_name": "deepseek-big",
        "litellm_params": {
            "model": "openai/gpt-4o-2024-05-13",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "5",
            "token_per_second": "24",
            "cost_per_token": "0.41",
        }
    },
    {
        "model_name": "deepseek-big",
        "litellm_params": {
            "model": "openai/gpt-4",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "6",
            "token_per_second": "32",
            "cost_per_token": "0.1",
        }
    },
    {
        "model_name": "deepseek-big",
        "litellm_params": {
            "model": "openai/gpt-4-0314",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "5",
            "token_per_second": "37",
            "cost_per_token": "0.11",
        }
    },
    {
        "model_name": "deepseek-small",
        "litellm_params": {
            "model": "openai/gpt-3.5-turbo",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "2",
            "token_per_second": "56",
            "cost_per_token": "0.12",
        }
    },
    {
        "model_name": "deepseek-small",
        "litellm_params": {
            "model": "openai/gpt-3.5-turbo",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "7",
            "token_per_second": "54",
            "cost_per_token": "0.11",
        }
    },
    {
        "model_name": "deepseek-small",
        "litellm_params": {
            "model": "openai/gpt-3.5-turbo-16k",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "2",
            "token_per_second": "24",
            "cost_per_token": "0.11",
        }
    },
    {
        "model_name": "deepseek-small",
        "litellm_params": {
            "model": "openai/o1-mini",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "6",
            "token_per_second": "37",
            "cost_per_token": "0.14",
        }
    },
    {
        "model_name": "deepseek-small",
        "litellm_params": {
            "model": "openai/gpt-4o-mini",
            "api_key": os.getenv("OPENAI_API_KEY"),
            "api_base": os.getenv("OPENAI_API_BASE"),
        },
        "model_info": {
            "first_token_time": "5",
            "token_per_second": "47",
            "cost_per_token": "0.41",
        }
    },
    
]   


# 创建路由策略实例
routing_strategy = CustomRoutingStrategy()
# 创建路由器并配置缓存
router = Router(model_list=model_list,
                cache_responses=True)
# 设置路由策略并传入router实例
routing_strategy.set_router(router)
router.set_custom_routing_strategy(routing_strategy)

print(dir(router))






async def test_simple_question():
    """测试简单问题"""
    print("\n=== 测试简单问题 ===")
    try:
        response = await router.acompletion(
            model="",  # 不指定具体模型，让路由策略自动选择
            messages=[{"role": "user", "content": "你好，请介绍同济大学"}]
        )
        print("响应:", response)
    except Exception as e:
        print(f"测试简单问题时出错: {str(e)}")


async def test_complex_question():
    """测试复杂问题"""
    print("\n=== 测试复杂问题 ===")
    try:
        response = await router.acompletion(
            model="",
            messages=[{
                "role": "user",
                "content": """请分析以下系统架构的优缺点，并给出优化建议：
                1. 前端使用aaaa
                2. 后端使用bbbbbbb
                3. 数据库使用cccccccc
                4. 使用ddddd做缓存
                5. 使用cccc做反向代理"""
            }]
        )
        print("响应:", response)
    except Exception as e:
        print(f"测试复杂问题时出错: {str(e)}")


async def main():
    """主测试函数"""
    try:


    # 测试简单问题
        await test_simple_question()
        await asyncio.sleep(2)

    # 测试复杂问题
        await test_complex_question()

    except Exception as e:
        print(f"测试过程中发生错误: {str(e)}")


if __name__ == "__main__":
    asyncio.run(main())